#!/usr/bin/env python
# coding=utf-8

import os
import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
from sklearn import preprocessing
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from sklearn.decomposition import PCA


def tsne_plotter(data, label, save_png, title):
    n_labels = len(set(label))

    # tsne
    tsne = TSNE(n_components=2, init='pca', learning_rate=10, perplexity=12, n_iter=1000)
    transformed_data = tsne.fit_transform(data)

    # LDA
    # lda = LDA(n_components=2)
    # transformed_data = lda.fit_transform(data, label)

    # PCA
    # pca = PCA(n_components=2)
    # transformed_data = pca.fit_transform(data)

    plt.figure()
    plt.scatter(transformed_data[:, 0], transformed_data[:, 1], 10, c=label, cmap=plt.cm.Spectral, alpha=0.5)
    plt.title(title)
    plt.legend()
    plt.savefig(save_png)

npy_dir = "dev"
spkers = os.listdir(npy_dir)
np.random.shuffle(spkers)

label = []
data = []
cnt = 0
print("data loading")
for spker in spkers:
    path = os.path.join(npy_dir, spker)
    npys = [ os.path.join(path, x) for x in os.listdir(path) ]
    for npy in npys:
        data.append(np.load(npy))
        label.append(cnt)
    cnt += 1
    print(cnt)

data = np.array(data)
label = np.array(label)

print("data shape: ", data.shape)
print("plotting...")
tsne_plotter(data=data, label=label, save_png="tsne.png", title="speaker embedding visualization")

